package com.radar.algorithm;

import cn.hutool.core.io.FileUtil;
import com.radar.utils.BitConverter;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.File;
import java.io.IOException;
import java.util.Arrays;
import java.util.stream.IntStream;

import javax.imageio.ImageIO;

/**
 * @author: fy-wbj
 * @description: Apply semi-Lagrangian backward extrapolation to a two-dimensional
 * precipitation field.
 * @date: 2020/12/30
 */
public class SemiLagrangian {

    public static void main(String[] args) throws IOException {
        byte[] bytes = FileUtil.readBytes("/Users/xx/data/fy4_flow.dat");
        byte[] bytes1 = FileUtil.readBytes("/Users/xx/data/precip.dat");
        float[][][] V = new float[2][1226][2001];
        float[][] precip = new float[1226][2001];
        int index = 0;
        for (int i = 0; i < V.length; i++) {
            for (int j = 0; j < V[0].length; j++) {
                for (int k = 0; k < V[0][0].length; k++) {
                    V[i][j][k] = (float) BitConverter.toDouble(bytes, index);
                    index += 8;
                }
            }
        }
        index = 0;
        for (int i = 0; i < precip.length; i++) {
            for (int j = 0; j < precip[0].length; j++) {
                precip[i][j] = (float) BitConverter.toDouble(bytes1, index);
                index += 8;
            }
        }
        SemiLagrangian semiLagrangian = new SemiLagrangian();
        float[][][] extrapolate = semiLagrangian.extrapolate(precip, V[0], V[1], 24);
        for (int i = 0; i < 24; i++) {
            BufferedImage bufferedImage = new BufferedImage(2001, 1226, BufferedImage.TYPE_INT_ARGB);
            for (int j = 0; j < 1226; j++) {
                for (int k = 0; k < 2001; k++) {
                    int z = (int) (extrapolate[i][j][k] == Float.NaN ? 0 : extrapolate[i][j][k]);
                    z = z > 255 ? 255 : z;
                    bufferedImage.setRGB(k, j, new Color(z, z, z).getRGB());
                }
            }
            ImageIO.write(bufferedImage, "png", new File("/Users/xx/data/java/flow_" + i + ".png"));
        }
    }

    float[][][] xy_coords = null;

    float[][][] XYW = null;

    float[][][] D = null;

    float[][][] V_inc = null;

    public float[][][] extrapolate(float[][] precip, float[][] u, float[][] v, int times) {
        int h = u.length;
        int w = u[0].length;
        float[][][] data = new float[times][h][w];
        D = new float[2][h][w];
        V_inc = new float[2][h][w];
        XYW = new float[2][h][w];
        for (int i = 0; i < 2; i++) {
            for (int j = 0; j < h; j++) {
                for (int k = 0; k < w; k++) {
                    V_inc[0][j][k] = u[j][k];
                    V_inc[1][j][k] = u[j][k];
                }
            }
        }
        xy_coords = new float[2][h][w];
        IntStream.range(0, h).forEach(z -> Arrays.fill(xy_coords[1][z], z));
        IntStream.range(0, w).forEach(z -> {
            for (int i = 0; i < h; i++) {
                xy_coords[0][i][z] = z;
            }
        });
        for (int i = 0; i < times; i++) {
            interpolateMotion(div(D, mul(V_inc, 0.5F)), 1, u, v);
            D = div(D, V_inc);
            interpolateMotion(D, 1, u, v);

            XYW = add(xy_coords, D);
            float[][] z = XYW[0];
            float[][] z1 = XYW[1];
            XYW[0] = z1;
            XYW[1] = z;
            data[i] = mapCoordinates(precip, XYW, true);
        }
        return data;
    }

    public void interpolateMotion(float[][][] D, int td, float[][] u, float[][] v) {
        XYW = add(xy_coords, D);
        float[][] z = XYW[0];
        float[][] z1 = XYW[1];
        XYW[0] = z1;
        XYW[1] = z;
        float[][] VWX = mapCoordinates(u, XYW);
        float[][] VWY = mapCoordinates(v, XYW);

        V_inc[0] = VWX;
        V_inc[1] = VWY;
        V_inc = mul(V_inc, td);
    }

    public float[][] mapCoordinates(float[][] precip, float[][][] XYW, boolean outval) {
        float[][] data = new float[precip.length][precip[0].length];
        for (int i = 0; i < precip.length; i++) {
            for (int j = 0; j < precip[0].length; j++) {
                int yindex = (int) Math.round(XYW[0][i][j]);
                int xindex = (int) Math.round(XYW[1][i][j]);
                if (yindex > 0 && yindex < precip.length - 1 && xindex > 0 && xindex < precip[0].length - 1) {
                    data[i][j] = precip[yindex][xindex];
                } else {
                    if (outval) {
                        data[i][j] = Float.NaN;
                    } else {
                        data[i][j] = precip[0][0];
                    }
                }
            }
        }
        return data;
    }

    public float[][] mapCoordinates(float[][] precip, float[][][] XYW) {
        return mapCoordinates(precip, XYW, false);
    }

    public float[][][] mul(float[][][] data, float v) {
        float[][][] c = new float[data.length][data[0].length][data[0][0].length];
        for (int i = 0; i < data.length; i++) {
            for (int j = 0; j < data[0].length; j++) {
                for (int k = 0; k < data[0][0].length; k++) {
                    c[i][j][k] = data[i][j][k] * v;
                }
            }
        }
        return c;
    }

    public float[][][] add(float[][][] data, float v) {
        float[][][] c = new float[data.length][data[0].length][data[0][0].length];
        for (int i = 0; i < data.length; i++) {
            for (int j = 0; j < data[0].length; j++) {
                for (int k = 0; k < data[0][0].length; k++) {
                    c[i][j][k] = data[i][j][k] + v;
                }
            }
        }
        return c;
    }

    public float[][][] add(float[][][] data, float[][][] data1) {
        float[][][] c = new float[data.length][data[0].length][data[0][0].length];
        for (int i = 0; i < data.length; i++) {
            for (int j = 0; j < data[0].length; j++) {
                for (int k = 0; k < data[0][0].length; k++) {
                    c[i][j][k] = data[i][j][k] + data1[i][j][k];
                }
            }
        }
        return c;
    }

    public float[][][] div(float[][][] data, float[][][] data1) {
        float[][][] c = new float[data.length][data[0].length][data[0][0].length];
        for (int i = 0; i < data.length; i++) {
            for (int j = 0; j < data[0].length; j++) {
                for (int k = 0; k < data[0][0].length; k++) {
                    c[i][j][k] = data[i][j][k] - data1[i][j][k];
                }
            }
        }
        return c;
    }
}
